"""
Mengji 调度器实现

MJ 调度器目的是自动化管理任务，但不是在任务执行流程，而是经可能让任务分阶段进行。
MJ 调度认为一个任务分为两个部分：**生成**阶段，这个阶段会评估需要进行哪些工作，将这些工作定义然后发送到任务队列
**运行**阶段，运行阶段是程序进行计算密集型任务的阶段，但是实际上 python 不完全适合执行计算密集型工作，所以在有的
地方会有计算后端这种方式出现，计算后端实际上是一种跨进程交互技术，将复杂的任务委托给其他进程处理，这些进程可以是本地进程，
或者是容器进程，还可以运行在其他主机上

date: 2025/3/19
author: SiHeng Tang
file: async_sche.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import asyncio
import gc
import heapq
import time
import logging

from mj_errors import MJRuntimeError

logger = logging.getLogger(__name__)


class ExitGenerationError(Exception):
    """
    为了解决封装带来的阶段控制和阶段过度嵌套的问题，使用生成退出类。
    生成异常与循环异常类似，是一类功能性异常，可以结束任务的生成流程，
    导致任务最终不会生成而是跳过插入任务的阶段，直接调用 ``_on_exit`` 函数清理资源
    """
    pass


class TaskWrapper(object):
    """
    实现跟踪和储存任务生命周期参数

    优先级
        优先级决定任务执行的顺序，在 ASche 的设计中，高于一个特定的优先级的任务一定会被执行,
        优先级的作用主要是决定任务执行的顺序，而非是抢占式系统中任务获取时间片的顺序。
        优先级的执行分界是 ``IDLE_PRIORITY`` 大于此优先级的任务会按照优先级从大到小排序，
        优先级越大的任务越靠前执行，同优先级任务执行顺序是不保证的。这里的特性被用来实现空闲计时器，
        和解决任务执行的顺序问题，可能在系统设计中有的任务必须优先执行。优先级可以解决 ASche 执行顺序问题。
    """
    IDLE_PRIORITY = -1

    def __init__(self, task, **kwargs):
        self.priority = kwargs["priority"] if "priority" in kwargs.keys() else self.IDLE_PRIORITY + 1
        self._task = task

    def __lt__(self, other):
        """
        重载比较函数，实现优先级执行顺序
        """
        # 注意 我们需要更大优先级项目优先执行
        return self.priority > other.priority

    def get_task(self):
        cache_t = self._task
        self._task = None
        return cache_t


class BaseTaskGen(object):
    """
    任务生成器有两种状态：等待和生成，生成器会不断在这两种状态之间循环。

    生成器的最终目标就是生成任务，一个生成器只对应一个任务。
    ``_on_gen`` 实现生成条件判断，``_func`` 作为任务的执行函数。

    为了满足内存管理需求，TG 不引用调度器，而是直接在生成流程中返回任务对象。
    生成器在任务执行中可能被卸载，这是为了满足可重入流程需求，卸载时调用 ``_on_exit``。

    继承基类，然后重写 ``_on_gen`` ``_run`` ``_on_exit``::

        class MyNode(BaseTaskGen):
            def __init__(self, sche):
                super().__init__(sche, wrapper_kwargs)
                # 根据需求重载基类
            def _on_gen(self):
                # 此处初始化函数执行，判断函数执行条件
            def _on_exit(self):
                # 资源回收，必须解除强引用
            async def _run(self):
                # 任务体
    """

    def __init__(self, **kwargs):
        """
        :param kwargs: 用于直接向任务包装器传递的参数
        """
        self._wrapper_kwargs = kwargs

    def generate_task(self):
        """
        生成逻辑：
            1. 调用 ``_on_gen`` 函数，判断资源状态，决策是否生成任务
            2. ``_on_gen`` 不终止生成，调用 ``_gen_task`` 生成任务并返回包装后对象
        """
        try:
            self._on_gen()
            return self._gen_task()
        except ExitGenerationError:
            return None

    def unload_generator(self):
        """
        卸载生成器，保证：
            1. 解除了对其他对象的引用
            2. 释放系统资源，例如网络连接，MAV 串行接口连接
        """
        # TODO 实现一个方法，可以主动结束和回收生成器
        self._on_exit()
        self._wrapper_kwargs = None

    async def _func(self):
        pass

    def _on_gen(self):
        pass

    def _on_exit(self):
        pass

    def _gen_task(self):
        return TaskWrapper(asyncio.create_task(self._func()), **self._wrapper_kwargs)


class TaskGenerator(BaseTaskGen):
    """
    更通用的任务生成器
    """

    def __init__(self, **kwargs):
        """
        通用任务生成器，在变长参数解析中新增：
            - ``interval`` 控制执行间隔，等效执行频率倒数
            - ``run`` 可执行异步函数
            - ``args`` 异步函数参数列表
            - ``kwargs`` 异步函数可变参数列表
        """
        super().__init__(priority=TaskWrapper.IDLE_PRIORITY + 10, **kwargs)
        self._func_kwargs = kwargs["kwargs"] if "kwargs" in kwargs.keys() else None
        self._func_args = kwargs["args"] if "args" in kwargs.keys() else None
        self._run_cro = kwargs["run"] if "run" in kwargs.keys() else None

        self._exec_interval = kwargs["interval"] if "interval" in kwargs.keys() else 0.05
        self._last_exec_time = time.time()

    def _on_gen(self):
        if (time.time() < (self._last_exec_time + self._exec_interval)) or (not bool(self._run_cro)):
            raise ExitGenerationError()
        self._last_exec_time = time.time()

    def _on_exit(self):
        self._func_kwargs = None
        self._func_args = None
        self._run_cro = None

    async def _func(self):
        if self._func_args and self._func_kwargs:
            await self._run_cro(*self._func_args, **self._func_kwargs)
        elif self._func_kwargs:
            await self._run_cro(**self._func_kwargs)
        elif self._func_args:
            await self._run_cro(*self._func_args)
        else:
            await self._run_cro()


class AsyncScheduler:
    """
    调度器，使用标准异步库实现
    """

    MAX_CYCLE_TIME = 0.01

    def __init__(self):
        self._generator = tuple()
        self._waiting_task = list()

        self._loop_start_t = time.time()
        self.idle_cnt = 0  # 空闲执行次数
        self.exec_cnt = 0  # 总执行次数
        self.cps_cnt = 0.  # 次每秒

    def query_status(self):
        return {"idle_cnt": self.idle_cnt,
                "exec_cnt": self.exec_cnt,
                "cps_cnt": self.cps_cnt,
                "loop_time": time.time() - self._loop_start_t,
                "static_gen_cnt": len(self._generator),
                "wait_tasks": len(self._waiting_task)}

    def reg_generator(self, g: BaseTaskGen):
        """
        :param g: 生成器对象
        """
        cache_list = list(self._generator)
        cache_list.append(g)
        self._generator = tuple(cache_list)

    def add_task(self, t: TaskWrapper):
        """
        :param t: 封装后的协程对象
        """
        heapq.heappush(self._waiting_task, t)

    def finalize(self):
        """
        调用所有生成器卸载函数，终止所有生成器生命周期。
        清空任务列表，清空生成器引用。
        """
        for g in self._generator:
            g.unload_generator()

        self._waiting_task.clear()
        self._generator = tuple()

    async def main_loop(self):
        while True:
            await self._main_loop()

    async def next_loop(self):
        await self._main_loop()

    async def _empty_task(self):
        """
        空任务实际上是一个优先级为 IDLE_PRIORITY 的任务。
        通过优先级控制，此任务在没有其他高优先级任务执行时执行，
        实现空闲时间管理，空任务每一次执行都会递增一次空闲计数器。
        """
        self.idle_cnt += 1
        await asyncio.sleep(0)

    async def _main_loop(self):
        """
        事件循环按照一定逻辑调度任务，事件循环被重复执行，事件循环结束调度器生命周期终止。
        事件循环一次按照以下顺序调度任务::
            1. 启动记录
                事件循环使用协程调度器循环获取启动时间戳，记录启动时 ``exec_cnt`` 执行次数。
            2. 轮询生成器
                事件循环轮询生成器，调用生成器的 ``generate_task`` 方法生成任务，
                最后调用 ``add_task`` 接口生成空任务。
            3. 轮询任务
                任务在生成时已经使用 ``heapq`` 堆排序，弹出任务执行，
                ``IDLE_PRIORITY`` 通过任务启动前比较优先级实现。
            4. 收集运行数据
                事件循环记录执行任务增量和执行时间实现系统性能估算。
        """
        start_time = time.time()
        start_exec_cnt = self.exec_cnt
        idle_exec_flag = 0  # 为 0 则执行空闲任务

        for g in self._generator:
            if bool(t := g.generate_task()):
                self.add_task(t)
        self.add_task(TaskWrapper(asyncio.create_task(self._empty_task()), priority=TaskWrapper.IDLE_PRIORITY))

        # TODO 加入安全崩溃流程，使用回调函数实现
        try:
            async with asyncio.timeout(self.MAX_CYCLE_TIME):
                for i in range(len(self._waiting_task)):
                    t = heapq.heappop(self._waiting_task)
                    if idle_exec_flag or t.priority > TaskWrapper.IDLE_PRIORITY:
                        idle_exec_flag = 1
                    if (not idle_exec_flag) or t.priority > TaskWrapper.IDLE_PRIORITY:
                        await t.get_task()
                        self.exec_cnt += 1

        except TimeoutError:
            logger.warning("TIME OUT! Is system overloaded or software closed loop?")

        finally:
            delta_time = time.time() - start_time
            delta_cnt = self.exec_cnt - start_exec_cnt
            self.cps_cnt = delta_cnt / delta_time
            self._waiting_task.clear()


MJScheduler = AsyncScheduler


# ################
# 任务调度器测试
# ################

async def print_time(_id):
    for i in range(10):
        print(f"My print id {_id=} Now time {time.time()=}")
    raise MJRuntimeError(0, "generator reentry frame test")


async def test_sche():
    while True:
        s = MJScheduler()

        gen_who = TaskGenerator(interval=1, run=print_time, display_del=1, args=("WHO SEC",))
        gen_hal = TaskGenerator(interval=0.5, run=print_time, display_del=1, args=("HALF SEC",))

        s.reg_generator(gen_who)
        s.reg_generator(gen_hal)
        try:
            await s.main_loop()
        except MJRuntimeError:
            s.finalize()

        gen_hal = None
        gen_who = None
        s = None

        gc.collect()
        # objgraph.show_backrefs(gc.garbage,filename="leak.png")
        print(f"gc {len(gc.garbage)}", gc.garbage)
        gc.garbage.clear()

        time.sleep(1)


if __name__ == '__main__':
    gc.set_debug(gc.DEBUG_LEAK)
    asyncio.run(test_sche())
